import casadi as ca
import numpy as np
from vehicle_model import BicycleModel

class MPCController:
    def __init__(self, dt=0.01, N=15, L=2.8):
        # MPC参数
        self.dt = dt  # 采样时间(s)
        self.N = N    # 预测时域长度
        self.L = L    # 轴距(m)
        self.vehicle_model = BicycleModel(wheelbase=L)

        # 状态和控制维度
        self.nx = self.vehicle_model.get_state_size()  # 状态维度: x, y, theta, v
        self.nu = self.vehicle_model.get_control_size()  # 控制维度: delta, a

        # 权重矩阵
        self.Q = np.diag([10.0, 10.0, 5.0, 1.0])  # 状态跟踪权重
        self.R = np.diag([1.0, 0.5])              # 控制输入权重
        self.Rd = np.diag([0.5, 0.2])             # 控制变化率权重

        # 控制约束
        self.delta_max = self.vehicle_model.max_steer  # 最大转向角(rad)
        self.a_max = self.vehicle_model.max_accel      # 最大加速度(m/s²)
        self.a_min = self.vehicle_model.max_decel      # 最大减速度(m/s²)
        self.delta_dot_max = 0.5                      # 最大转向角速度(rad/s)
        self.a_dot_max = 2.0                           # 最大加速度变化率(m/s³)

        # 初始化优化问题
        self.setup_mpc_problem()

    def setup_mpc_problem(self):
        # 创建符号变量
        self.x = ca.SX.sym('x', self.nx)
        self.u = ca.SX.sym('u', self.nu)
        self.x_ref = ca.SX.sym('x_ref', self.nx)

        # 定义状态方程 (使用自行车模型)
        delta, a = self.u[0], self.u[1]
        x_dot = ca.vertcat(
            self.x[3] * ca.cos(self.x[2]),
            self.x[3] * ca.sin(self.x[2]),
            self.x[3] / self.L * ca.tan(delta),
            a
        )

        # 离散化状态方程 (欧拉法)
        self.f = ca.Function('f', [self.x, self.u], [self.x + self.dt * x_dot])

        # 优化变量
        self.opt_x = ca.SX.sym('opt_x', self.nx, self.N+1)
        self.opt_u = ca.SX.sym('opt_u', self.nu, self.N)

        # 代价函数
        self.cost = 0
        for k in range(self.N):
            # 跟踪误差代价
            x_err = self.opt_x[:, k] - self.x_ref
            self.cost += ca.mtimes([x_err.T, self.Q, x_err])
            # 控制输入代价
            self.cost += ca.mtimes([self.opt_u[:, k].T, self.R, self.opt_u[:, k]])

            # 控制变化率代价
            if k > 0:
                du = self.opt_u[:, k] - self.opt_u[:, k-1]
                self.cost += ca.mtimes([du.T, self.Rd, du])

        # 约束条件
        self.constraints = []
        self.lb = []
        self.ub = []

        # 初始状态约束
        self.constraints += [self.opt_x[:, 0]]
        self.lb += [np.zeros(self.nx)]
        self.ub += [np.zeros(self.nx)]

        # 状态和控制约束
        for k in range(self.N):
            # 状态约束 (车辆动力学)
            x_next = self.f(self.opt_x[:, k], self.opt_u[:, k])
            self.constraints += [self.opt_x[:, k+1] - x_next]
            self.lb += [np.zeros(self.nx)]
            self.ub += [np.zeros(self.nx)]

            # 控制输入约束
            self.constraints += [self.opt_u[0, k]]  # delta
            self.lb += [-self.delta_max]
            self.ub += [self.delta_max]

            self.constraints += [self.opt_u[1, k]]  # a
            self.lb += [self.a_min]
            self.ub += [self.a_max]

            # 控制变化率约束
            if k > 0:
                self.constraints += [self.opt_u[0, k] - self.opt_u[0, k-1]]
                self.lb += [-self.delta_dot_max * self.dt]
                self.ub += [self.delta_dot_max * self.dt]

                self.constraints += [self.opt_u[1, k] - self.opt_u[1, k-1]]
                self.lb += [-self.a_dot_max * self.dt]
                self.ub += [self.a_dot_max * self.dt]

        # 创建优化问题
        self.opti = ca.Opti()
        self.opti.minimize(self.cost)
        self.opti.subject_to(self.constraints)

        # 设置求解器
        opts = {
            'ipopt': {
                'max_iter': 100,
                'print_level': 0,
                'acceptable_tol': 1e-6,
                'acceptable_obj_change_tol': 1e-4
            },
            'print_time': 0
        }
        self.opti.solver('ipopt', opts)

    def compute_control(self, x0, x_ref, prev_u=None):
        # 设置初始状态
        self.opti.set_value(self.opt_x[:, 0], x0)

        # 设置参考轨迹
        self.opti.set_value(self.x_ref, x_ref)

        # 设置初始猜测 (温启动)
        if prev_u is not None:
            for k in range(min(self.N, len(prev_u))):
                self.opti.set_initial(self.opt_u[:, k], prev_u[k])

        # 求解优化问题
        try:
            sol = self.opti.solve()
            u_opt = sol.value(self.opt_u)
            return u_opt[:, 0].reshape(-1, 1)
        except:
            # 优化失败时返回零控制
            return np.zeros((self.nu, 1))

    def get_prediction_horizon(self):
        return self.N

    def get_sampling_time(self):
        return self.dt